#PBAR_Zspec_GenerateCalibration.py

import PBAR_Zspec
import numpy as np
import os
import csv
# Location of data
basepath = r'C:\Users\jkwong\Documents\Work\PBAR\data'

#countRange = np.array([0.7e3, 1e3])
countRange = np.array([250, 325])

# range in which the spectra mush intersect with the horizontal range
binBounds = np.array([20, 200])

# detector to match
detStandard = 69

# list of good/bad detectors
(goodDetectorsList, badDetectorsList) = \
                    PBAR_Zspec.GenerateDefaultDetectorList()
# Load zspec data
filenameList = ['ea26.csv', 'ea27.csv', 'ea49.csv', 'ea50.csv', 'ea51.csv']

fullfilenameList = []
for i in xrange(len(filenameList)):
    fullfilenameList.append(os.path.join(basepath, filenameList[i]))

dat = PBAR_Zspec.ReadZspec(fullfilenameList)

## Load summary data
#infoFilename = os.path.join(basepath, 'datasetSummary.txt')
#(datasetDescription, datasetAcquisitionTime, \
#    datasetTime, datasetTimeNum, datasetTimeStr) = \
#    PBAR_Zspec.GetDatasetInformation(infoFilename, filenameList)    


# Shift among PMTs
ii = 1
pmtShift = np.zeros(dat[0].shape[1])
#bins = np.arange(dat[0].shape[0])
binArray = np.matlib.repmat(np.arange(0,dat[0].shape[0]),dat[0].shape[1],1).T
cutt = (binArray > binBounds[0]) & (binArray < binBounds[1]) & (dat[ii] > countRange[0]) & (dat[ii] < countRange[1])
binMean = (binArray * dat[ii] * cutt).sum(axis = 0).astype(float) / (dat[ii] * cutt).sum(axis = 0).astype(float) # min bin, calculated using the count as weight
binMean = binMean.astype(float)
pmtShift = binMean[detStandard] / binMean
pmtShift[np.isnan(pmtShift)] = -1


# output the values to file
fileOutput = os.path.join(basepath, 'gainValues20130816_1057.txt')
fid = open(fileOutput, 'wb')
csvWriterObj = csv.writer(fid, delimiter = ' ')
for i in xrange(len(pmtShift)):
    csvWriterObj.writerow(np.array([i, '%2.5f' %pmtShift[i]]))
fid.close()



# Calculate and output the counts
index = 1
threshold = 14
counts = np.zeros(dat[index].shape[1])
for i in xrange(len(counts)):
    x = np.arange(256) * pmtShift[i]
    cut = x >=14
    counts[i] = dat[index][cut,i].sum()
fileOutput = os.path.join(basepath, 'counts_threshold14_20130816_1057.txt')
fid = open(fileOutput, 'wb')
csvWriterObj = csv.writer(fid, delimiter = ' ')
for i in xrange(len(counts)):
    csvWriterObj.writerow(np.array([i, '%d' %counts[i]]))
fid.close()



#
## measure for calculating the gain
#binMeanCC = np.zeros((len(dat), dat[0].shape[1]))
## gain shift across datasets
#gainShift = np.zeros((len(dat), dat[0].shape[1]))
#
#binArray = np.matlib.repmat(np.arange(0,dat[0].shape[0]),dat[0].shape[1],1).T
#
#for cc in range(len(datasetGroupsIndices['CC'])):
#    ii = CCindices[cc]
#    cutt = (binArray > binBounds[0]) & (binArray < binBounds[1]) & (dat[ii] > countRange[0]) & (dat[ii] < countRange[1])
#    binMeanCC[cc,:] = (binArray * dat[ii] * cutt).sum(axis = 0) / (dat[ii] * cutt).sum(axis = 0)
#    gainShift[cc,:] = binMeanCC[0,:] / binMeanCC[cc,:]
#
## Build the array that will be written to file, containing time string, epoch time and gain correction values
#temp = np.empty((len(CCindices), 1 + dat[0].shape[1]), dtype = '|S20')
#temp[:,0] = datasetTimeNum[CCindices]
#temp[:,1:] = binMeanCC
##temp[:,1:] = gainShift



#########################
# EXPLORATORY PLOTS
#########################

index = 1

plt.figure()
plt.grid()

listlist = copy.copy(goodDetectorsList)
listlist = copy.copy(badDetectorsList)
listlist = np.arange(dat[index].shape[1])


for i in xrange(len(listlist)):
    detIndex = listlist[i]
#for i in xrange(136):
#    detIndex = i
    plt.plot(np.arange(256), dat[index][:,detIndex], label = 'Det# %d' %detIndex)
    
# plot the count range
plot([0, 255], countRange[0] * np.ones(2), '--b')
plot([0, 255], countRange[1] * np.ones(2), '--b')

plt.yscale('log')
#plt.legend()
plt.xlim([0, 255])

plt.title('%s' % filenameList[index])
plt.xlabel('Bin')
plt.ylabel('Counts')


# with gain correction

index = 1

plt.figure()
plt.grid()

listlist = copy.copy(goodDetectorsList)
listlist = copy.copy(badDetectorsList)
listlist = np.arange(dat[index].shape[1])


for i in xrange(len(listlist)):
    detIndex = listlist[i]
#for i in xrange(136):
#    detIndex = i
    plt.plot(np.arange(256) * pmtShift[detIndex], dat[index][:,detIndex], label = 'Det# %d' %detIndex)
    
# plot the count range
plot([0, 255], countRange[0] * np.ones(2), '--b')
plot([0, 255], countRange[1] * np.ones(2), '--b')

plot(np.array([14, 14]), np.array([1, 1e6]), '-b')

plt.yscale('log')
#plt.legend()
plt.xlim([0, 255])

plt.title('%s' % filenameList[index])
plt.xlabel('Bin')
plt.ylabel('Counts')



# Bin value versus detector number

plt.figure()
plt.grid()
plot(np.arange(136), binMean, '*k')
plt.ylabel('Bin')
plt.xlabel('Detector')


# Bin value versus detector number

plt.figure()
plt.grid()
plot(np.arange(136), counts, '*k')
plt.ylabel('Counts')
plt.xlabel('Detector')

